import glob
import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
from tqdm import tqdm

from preprocessing.options import parse_args


def extract_train(param, args, crop_size=128):
    image_path, mask_path = param
    img_dir = os.path.join(args.root_dir, args.crops_dir, Path(image_path).name[:-4])
    img_mask_dir = os.path.join(args.root_dir, args.mask_crops_dir, Path(image_path).name[:-4])
    os.makedirs(img_dir, exist_ok=True)
    os.makedirs(img_mask_dir, exist_ok=True)

    img = cv2.imread(image_path, cv2.IMREAD_COLOR)
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    assert img.shape[0] == mask.shape[0] and img.shape[1] == mask.shape[1]
    height, width = img.shape[0], img.shape[1]
    for row in range(height // crop_size + 1):
        for col in range(width // crop_size + 1):
            row_end = min(height, (row + 1) * crop_size)
            row_start = row_end - crop_size
            col_end = min(width, (col + 1) * crop_size)
            col_start = col_end - crop_size

            if (row == height // crop_size and row_start == height) or \
                    (col == width // crop_size and col_start == width):
                continue

            img_crop_path = os.path.join(img_dir, "{}_{}.png".format(row, col))
            if os.path.exists(img_crop_path):
                continue
            img_crop = img[row_start:row_end, col_start:col_end]
            cv2.imwrite(img_crop_path, img_crop)

            mask_crop_path = os.path.join(img_mask_dir, "{}_{}.png".format(row, col))
            if os.path.exists(mask_crop_path):
                continue
            mask_crop = mask[row_start:row_end, col_start:col_end]
            cv2.imwrite(mask_crop_path, mask_crop)


def extract_test(param, args, crop_size=128):
    image_path = param
    img_dir = os.path.join(args.root_dir, args.crops_dir, Path(image_path).name[:-4])
    os.makedirs(img_dir, exist_ok=True)

    img = cv2.imread(image_path, cv2.IMREAD_COLOR)
    assert img.shape[1] == img.shape[1]
    height, width = img.shape[0], img.shape[1]
    for row in range(height // crop_size + 1):
        for col in range(width // crop_size + 1):
            row_end = min(height, (row + 1) * crop_size)
            row_start = row_end - crop_size
            col_end = min(width, (col + 1) * crop_size)
            col_start = col_end - crop_size

            if (row == height // crop_size and row_start == height) or \
                    (col == width // crop_size and col_start == width):
                continue

            img_crop_path = os.path.join(img_dir, "{}_{}.png".format(row, col))
            if os.path.exists(img_crop_path):
                continue
            img_crop = img[row_start:row_end, col_start:col_end]
            cv2.imwrite(img_crop_path, img_crop)


def get_image_mask_pairs(train_path, train_mask_path):
    image_mask_pairs = []
    for image_path in glob.glob(os.path.join(train_path, '*.jpg')):
        mask_path = os.path.join(train_mask_path, '{}.png'.format(Path(image_path).name[:-4]))
        image_mask_pairs.append((image_path, mask_path))
    return image_mask_pairs


def main():
    args = parse_args()
    train_path = os.path.join(args.root_dir, 'train')
    train_mask_path = os.path.join(args.root_dir, 'train_mask')
    train_image_mask_pairs = get_image_mask_pairs(train_path, train_mask_path)
    # with Pool(processes=1) as p:
    #     with tqdm(total=len(train_image_mask_pairs)) as pbar:
    #         func = partial(extract_train, args=args)
    #         for _ in p.imap_unordered(func, train_image_mask_pairs):
    #             pbar.update()

    test_path = os.path.join(args.root_dir, 'test')
    test_images = glob.glob(os.path.join(test_path, '*.jpg'))
    args.crops_dir = 'test_crops'
    with Pool(processes=1) as p:
        with tqdm(total=len(test_images)) as pbar:
            func = partial(extract_test, args=args)
            for _ in p.imap_unordered(func, test_images):
                pbar.update()


if __name__ == '__main__':
    main()
